package org.practice.springaichatmemory.controller;

import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import org.practice.springaichatmemory.service.ChatCleanupService;
import org.practice.springaichatmemory.service.SystemPromptService;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@RestController
public class ChatController {

    private final OpenAiChatModel chatModel;
    private final ChatMemory chatMemory;
    private final ChatCleanupService chatCleanupService;
    private final SystemPromptService systemPromptService;
    private final ExecutorService executor = Executors.newSingleThreadExecutor();

    @Autowired
    public ChatController(OpenAiChatModel chatModel, ChatMemory chatMemory, 
                         ChatCleanupService chatCleanupService, SystemPromptService systemPromptService) {
        this.chatModel = chatModel;
        this.chatMemory = chatMemory;
        this.chatCleanupService = chatCleanupService;
        this.systemPromptService = systemPromptService;
    }

    @GetMapping("/chat")
    public String chat(@RequestParam(value = "message", defaultValue = "Hello") String message,
                       @RequestParam(value = "sessionId", defaultValue = "default") String sessionId) {
        // 更新会话访问时间
        chatCleanupService.updateSessionAccessTime(sessionId);
        
        // 为会话初始化system prompt（如果需要）
        systemPromptService.initializeSystemPromptIfNeeded(sessionId);
        
        // 将用户消息添加到历史记录中
        chatMemory.add(sessionId, new UserMessage(message));
        
        // 获取所有历史消息（从配置文件中读取最大消息数限制）
        var history = chatMemory.get(sessionId, chatCleanupService.getMaxMessagesPerSession());
        
        // 创建包含历史记录的提示
        Prompt prompt = new Prompt(history);
        String response = chatModel.call(prompt).getResult().getOutput().toString();
        
        // 将AI响应也添加到历史记录中
        chatMemory.add(sessionId, new org.springframework.ai.chat.messages.AssistantMessage(response));

        System.out.println("Response: " + response);
        return response;
    }
    
    @GetMapping(value = "/chat-stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter chatStream(@RequestParam(value = "message", defaultValue = "Hello") String message,
                                 @RequestParam(value = "sessionId", defaultValue = "default") String sessionId) {
        // 创建SseEmitter用于流式响应
        SseEmitter emitter = new SseEmitter(Long.MAX_VALUE);
        
        // 在单独的线程中处理流式响应
        executor.submit(() -> {
            try {
                // 更新会话访问时间
                chatCleanupService.updateSessionAccessTime(sessionId);
                
                // 为会话初始化system prompt（如果需要）
                systemPromptService.initializeSystemPromptIfNeeded(sessionId);
                
                // 将用户消息添加到历史记录中
                chatMemory.add(sessionId, new UserMessage(message));
                
                // 获取所有历史消息（从配置文件中读取最大消息数限制）
                var history = chatMemory.get(sessionId, chatCleanupService.getMaxMessagesPerSession());
                
                // 创建包含历史记录的提示
                Prompt prompt = new Prompt(history);
                
                // 获取流式响应
                Flux<org.springframework.ai.chat.model.ChatResponse> stream = chatModel.stream(prompt);
                
                // 收集完整的响应以保存到历史记录中
                StringBuilder fullResponse = new StringBuilder();
                
                // 订阅流并发送每个块到客户端
                stream.subscribe(
                    chatResponse -> {
                        try {
                            String content = chatResponse.getResult().getOutput().toString();
                            fullResponse.append(content);
                            emitter.send(content);
                        } catch (IOException e) {
                            emitter.completeWithError(e);
                        }
                    },
                        emitter::completeWithError,
                    () -> {
                        try {
                            // 流完成时，将完整的响应添加到历史记录中
                            chatMemory.add(sessionId, new org.springframework.ai.chat.messages.AssistantMessage(fullResponse.toString()));
                            emitter.complete();
                        } catch (Exception e) {
                            emitter.completeWithError(e);
                        }
                    }
                );
            } catch (Exception e) {
                try {
                    emitter.completeWithError(e);
                } catch (Exception ex) {
                    // 忽略异常
                }
            }
        });
        
        return emitter;
    }
    
    @GetMapping("/history")
    public java.util.List<org.springframework.ai.chat.messages.Message> getHistory(
            @RequestParam(value = "sessionId", defaultValue = "default") String sessionId) {
        // 更新会话访问时间
        chatCleanupService.updateSessionAccessTime(sessionId);
        // 获取历史消息（从配置文件中读取最大消息数限制）
        return chatMemory.get(sessionId, chatCleanupService.getMaxMessagesPerSession());
    }
    
    @GetMapping("/clear")
    public String clearHistory(
            @RequestParam(value = "sessionId", defaultValue = "default") String sessionId) {
        chatMemory.clear(sessionId);
        // 清除会话的system prompt初始化状态
        systemPromptService.clearSessionInitialization(sessionId);
        return "History cleared for session: " + sessionId;
    }
    
    @GetMapping("/cleanup")
    public String cleanupExpiredSessions() {
        chatCleanupService.cleanupExpiredSessions();
        return "Cleanup task completed. Tracked sessions: " + chatCleanupService.getTrackedSessionCount();
    }
}